with Ada.Text_IO;  use Ada.Text_IO;
with Ada.Integer_Text_IO; use Ada.Integer_Text_IO;

package body AES
is
   package M_IO is new Ada.Text_IO.Modular_IO(Byte);
   
   procedure Print_Word(Arr: in Word) is
   begin
      for Index in Arr'Range loop
         M_IO.Put(Arr(Index));
         Put(" ");
      end loop;
   end;
   
   function mul(a, b : byte) return byte
   is
      r: byte;
   begin
      if a /= 0 and b /= 0 then
         r := Alogtable(byte((Integer(Logtable(a)) + 
                                Integer(Logtable(b))) mod 255));
      else
         r := 0;
      end if;
      
      --  New_Line;
      --  Put_Line("mul: a b r");
      --  M_IO.put(a); Put(" | ");
      --  M_IO.Put(b); Put(" | ");
      --  M_IO.Put(r); Put(" | ");
      --  New_Line;

      return r;
   end mul;

   function word_xor(a, b: word) return word
   is
      r: Word;
   begin
      r := word'(a(0) xor b(0), a(1) xor b(1), a(2) xor b(2), a(3) xor b(3));
      --  Put_Line("wordxXor: a b r");
      --  Print_Word(a); Put(" | ");
      --  Print_Word(b); Put(" | ");
      --  Print_Word(r); Put(" | ");
      --  New_Line;
      return r;
   end word_xor;

   function word_xor_xor(a, b, c: word) return word
   is
   begin
      return word'(a(0) xor b(0) xor c(0),
                   a(1) xor b(1) xor c(1),
                   a(2) xor b(2) xor c(2),
                   a(3) xor b(3) xor c(3));
   end word_xor_xor;


   function SubWord(w: word) return word
   is
   begin
      return word'(S(w(0)), S(w(1)), S(w(2)), S(w(3)));
   end SubWord;

   function RotWord(w: word) return word
   is
   begin
      return word'(w(1), w(2), w(3), w(0));
   end RotWord;

   function Block2State(t: block) return state
   is
   begin
      return state'(word'(t( 0), t( 1), t( 2), t( 3)),
                    word'(t( 4), t( 5), t( 6), t( 7)),
                    word'(t( 8), t( 9), t(10), t(11)),
                    word'(t(12), t(13), t(14), t(15)));
   end Block2State;

   function State2Block(st: state) return block
   is
   begin
      return block'(st(0)(0), st(0)(1), st(0)(2), st(0)(3),
                    st(1)(0), st(1)(1), st(1)(2), st(1)(3),
                    st(2)(0), st(2)(1), st(2)(2), st(2)(3),
                    st(3)(0), st(3)(1), st(3)(2), st(3)(3));
   end State2Block;

   function SubBytes(st: state) return state
   is
   begin
      return state'(word'(S(st(0)(0)), S(st(0)(1)), S(st(0)(2)), S(st(0)(3))),
                    word'(S(st(1)(0)), S(st(1)(1)), S(st(1)(2)), S(st(1)(3))),
                    word'(S(st(2)(0)), S(st(2)(1)), S(st(2)(2)), S(st(2)(3))),
                    word'(S(st(3)(0)), S(st(3)(1)), S(st(3)(2)), S(st(3)(3))));
   end SubBytes;

   function InvSubBytes(st: state) return state
   is
   begin
      return state'(word'(Si(st(0)(0)), Si(st(0)(1)), Si(st(0)(2)), Si(st(0)(3))),
                    word'(Si(st(1)(0)), Si(st(1)(1)), Si(st(1)(2)), Si(st(1)(3))),
                    word'(Si(st(2)(0)), Si(st(2)(1)), Si(st(2)(2)), Si(st(2)(3))),
                    word'(Si(st(3)(0)), Si(st(3)(1)), Si(st(3)(2)), Si(st(3)(3))));
   end InvSubBytes;

   function ShiftRows(st: state) return state
   is
   begin
      return state'(word'(st(0)(0), st(1)(1), st(2)(2), st(3)(3)),
                    word'(st(1)(0), st(2)(1), st(3)(2), st(0)(3)),
                    word'(st(2)(0), st(3)(1), st(0)(2), st(1)(3)),
                    word'(st(3)(0), st(0)(1), st(1)(2), st(2)(3)));
   end ShiftRows;

   function InvShiftRows(st: state) return state
   is
   begin
      return state'(word'(st(0)(0), st(3)(1), st(2)(2), st(1)(3)),
                    word'(st(1)(0), st(0)(1), st(3)(2), st(2)(3)),
                    word'(st(2)(0), st(1)(1), st(0)(2), st(3)(3)),
                    word'(st(3)(0), st(2)(1), st(1)(2), st(0)(3)));
   end InvShiftRows;

   function MixColumns(st: state) return state
   is
   begin
      return state'(
                    word'(
                          mul(2, st(0)(0)) xor mul(3, st(0)(1)) xor        st(0)(2)  xor        st(0)(3) ,
                          st(0)(0)  xor mul(2, st(0)(1)) xor mul(3, st(0)(2)) xor        st(0)(3) ,
                          st(0)(0)  xor        st(0)(1)  xor mul(2, st(0)(2)) xor mul(3, st(0)(3)),
                          mul(3, st(0)(0)) xor        st(0)(1)  xor        st(0)(2)  xor mul(2, st(0)(3))
                         ),
                    word'(
                          mul(2, st(1)(0)) xor mul(3, st(1)(1)) xor        st(1)(2)  xor        st(1)(3) ,
                          st(1)(0)  xor mul(2, st(1)(1)) xor mul(3, st(1)(2)) xor        st(1)(3) ,
                          st(1)(0)  xor        st(1)(1)  xor mul(2, st(1)(2)) xor mul(3, st(1)(3)),
                          mul(3, st(1)(0)) xor        st(1)(1)  xor        st(1)(2)  xor mul(2, st(1)(3))
                         ),
                    word'(
                          mul(2, st(2)(0)) xor mul(3, st(2)(1)) xor        st(2)(2)  xor        st(2)(3) ,
                          st(2)(0)  xor mul(2, st(2)(1)) xor mul(3, st(2)(2)) xor        st(2)(3) ,
                          st(2)(0)  xor        st(2)(1)  xor mul(2, st(2)(2)) xor mul(3, st(2)(3)),
                          mul(3, st(2)(0)) xor        st(2)(1)  xor        st(2)(2)  xor mul(2, st(2)(3))
                         ),
                    word'(
                          mul(2, st(3)(0)) xor mul(3, st(3)(1)) xor        st(3)(2)  xor        st(3)(3) ,
                          st(3)(0)  xor mul(2, st(3)(1)) xor mul(3, st(3)(2)) xor        st(3)(3) ,
                          st(3)(0)  xor        st(3)(1)  xor mul(2, st(3)(2)) xor mul(3, st(3)(3)),
                          mul(3, st(3)(0)) xor        st(3)(1)  xor        st(3)(2)  xor mul(2, st(3)(3))
                         )
                   );
   end MixColumns;

   function InvMixColumns(st: state) return state
   is
   begin
      return state'(
                    word'(
                          mul(16#0e#, st(0)(0)) xor mul(16#0b#, st(0)(1)) xor mul(16#0d#, st(0)(2)) xor mul(16#09#, st(0)(3)),
                          mul(16#09#, st(0)(0)) xor mul(16#0e#, st(0)(1)) xor mul(16#0b#, st(0)(2)) xor mul(16#0d#, st(0)(3)),
                          mul(16#0d#, st(0)(0)) xor mul(16#09#, st(0)(1)) xor mul(16#0e#, st(0)(2)) xor mul(16#0b#, st(0)(3)),
                          mul(16#0b#, st(0)(0)) xor mul(16#0d#, st(0)(1)) xor mul(16#09#, st(0)(2)) xor mul(16#0e#, st(0)(3))
                         ),
                    word'(
                          mul(16#0e#, st(1)(0)) xor mul(16#0b#, st(1)(1)) xor mul(16#0d#, st(1)(2)) xor mul(16#09#, st(1)(3)),
                          mul(16#09#, st(1)(0)) xor mul(16#0e#, st(1)(1)) xor mul(16#0b#, st(1)(2)) xor mul(16#0d#, st(1)(3)),
                          mul(16#0d#, st(1)(0)) xor mul(16#09#, st(1)(1)) xor mul(16#0e#, st(1)(2)) xor mul(16#0b#, st(1)(3)),
                          mul(16#0b#, st(1)(0)) xor mul(16#0d#, st(1)(1)) xor mul(16#09#, st(1)(2)) xor mul(16#0e#, st(1)(3))
                         ),
                    word'(
                          mul(16#0e#, st(2)(0)) xor mul(16#0b#, st(2)(1)) xor mul(16#0d#, st(2)(2)) xor mul(16#09#, st(2)(3)),
                          mul(16#09#, st(2)(0)) xor mul(16#0e#, st(2)(1)) xor mul(16#0b#, st(2)(2)) xor mul(16#0d#, st(2)(3)),
                          mul(16#0d#, st(2)(0)) xor mul(16#09#, st(2)(1)) xor mul(16#0e#, st(2)(2)) xor mul(16#0b#, st(2)(3)),
                          mul(16#0b#, st(2)(0)) xor mul(16#0d#, st(2)(1)) xor mul(16#09#, st(2)(2)) xor mul(16#0e#, st(2)(3))
                         ),
                    word'(
                          mul(16#0e#, st(3)(0)) xor mul(16#0b#, st(3)(1)) xor mul(16#0d#, st(3)(2)) xor mul(16#09#, st(3)(3)),
                          mul(16#09#, st(3)(0)) xor mul(16#0e#, st(3)(1)) xor mul(16#0b#, st(3)(2)) xor mul(16#0d#, st(3)(3)),
                          mul(16#0d#, st(3)(0)) xor mul(16#09#, st(3)(1)) xor mul(16#0e#, st(3)(2)) xor mul(16#0b#, st(3)(3)),
                          mul(16#0b#, st(3)(0)) xor mul(16#0d#, st(3)(1)) xor mul(16#09#, st(3)(2)) xor mul(16#0e#, st(3)(3))
                         )
                   );
   end InvMixColumns;

   function AddRoundKey(st: state; rk0, rk1, rk2, rk3: in word) return state
   is
   begin
      return state'(
                    word'(st(0)(0) xor rk0(0), st(0)(1) xor rk0(1), st(0)(2) xor rk0(2), st(0)(3) xor rk0(3)),
                    word'(st(1)(0) xor rk1(0), st(1)(1) xor rk1(1), st(1)(2) xor rk1(2), st(1)(3) xor rk1(3)),
                    word'(st(2)(0) xor rk2(0), st(2)(1) xor rk2(1), st(2)(2) xor rk2(2), st(2)(3) xor rk2(3)),
                    word'(st(3)(0) xor rk3(0), st(3)(1) xor rk3(1), st(3)(2) xor rk3(2), st(3)(3) xor rk3(3))
                   );
   end AddRoundKey;

   function KeySetupEnc4(cipherKey: key) return key_schedule
   is
      rk: key_schedule;
   begin
      rk := key_schedule'(others => word'(0, 0, 0, 0));

      for i in Integer range 0 .. 3 loop
         rk(i) := word'(cipherKey(4*i), cipherKey(4*i+1), cipherKey(4*i+2), cipherKey(4*i+3));
         --# assert 0 <= i and i <= 3 and
         --#        (for all p in Integer range 0 .. i =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3))));
      end loop;

      for i in Integer range 4 .. 43 loop
         if (i mod 4 = 0) then
            rk(i) := word_xor_xor(rk(i-4), SubWord(RotWord(rk(i-1))), rcon(i/4-1));
         else
            rk(i) := word_xor(rk(i-4), rk(i-1));
         end if;
         --# assert 4 <= i and i <= 43 and
         --#        (for all p in Integer range 0 .. 3 =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3)))) and
         --#        (for all p in Integer range 4 .. i =>
         --#          ((p mod 4 = 0 and rk(p) = word_xor_xor(rk(p-4), SubWord(RotWord(rk(p-1))), rcon(p/4-1))) or
         --#          (p mod 4 /= 0 and rk(p) = word_xor(rk(p-4), rk(p-1)))));
      end loop;

      return rk;
   end KeySetupEnc4;

   function KeySetupEnc6(cipherKey: key) return key_schedule
   is
      rk: key_schedule;
   begin
      rk := key_schedule'(others => word'(0, 0, 0, 0));

      for i in Integer range 0 .. 5 loop
         rk(i) := word'(cipherKey(4*i), cipherKey(4*i+1), cipherKey(4*i+2), cipherKey(4*i+3));
         --# assert 0 <= i and i <= 5 and
         --#        (for all p in Integer range 0 .. i =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3))));
      end loop;

      for i in Integer range 6 .. 51 loop
         if (i mod 6 = 0) then
            rk(i) := word_xor_xor(rk(i-6), SubWord(RotWord(rk(i-1))), rcon(i/6-1));
         else
            rk(i) := word_xor(rk(i-6), rk(i-1));
         end if;
         --# assert 6 <= i and i <= 51 and
         --#        (for all p in Integer range 0 .. 5 =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3)))) and
         --#        (for all p in Integer range 6 .. i =>
         --#          ((p mod 6 = 0 and rk(p) = word_xor_xor(rk(p-6), SubWord(RotWord(rk(p-1))), rcon(p/6-1))) or
         --#          (p mod 6 /= 0 and rk(p) = word_xor(rk(p-6), rk(p-1)))));
      end loop;

      return rk;
   end KeySetupEnc6;

   function KeySetupEnc8(cipherKey: key) return key_schedule
   is
      rk: key_schedule;
   begin
      rk := key_schedule'(others => word'(0, 0, 0, 0));

      for i in Integer range 0 .. 7 loop
         rk(i) := word'(cipherKey(4*i), cipherKey(4*i+1), cipherKey(4*i+2), cipherKey(4*i+3));
         --# assert 0 <= i and i <= 7 and
         --#        (for all p in Integer range 0 .. i =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3))));
      end loop;

      for i in Integer range 8 .. 59 loop
         if (i mod 8 = 0) then
            rk(i) := word_xor_xor(rk(i-8), SubWord(RotWord(rk(i-1))), rcon(i/8-1));
         elsif (i mod 4 = 0) then
            rk(i) := word_xor(rk(i-8), SubWord(rk(i-1)));
         else
            rk(i) := word_xor(rk(i-8), rk(i-1));
         end if;
         --# assert 8 <= i and i <= 59 and
         --#        (for all p in Integer range 0 .. 7 =>
         --#          (rk(p) = word'(cipherKey(4*p), cipherKey(4*p+1), cipherKey(4*p+2), cipherKey(4*p+3)))) and
         --#        (for all p in Integer range 8 .. i =>
         --#          ((p mod 8 = 0 and                  rk(p) = word_xor_xor(rk(p-8), SubWord(RotWord(rk(p-1))), rcon(p/8-1))) or
         --#          (p mod 8 /= 0 and p mod 4 = 0 and  rk(p) = word_xor(rk(p-8), SubWord(rk(p-1)))) or
         --#          (p mod 8 /= 0 and p mod 4 /= 0 and rk(p) = word_xor(rk(p-8), rk(p-1)))));
      end loop;

      return rk;
   end KeySetupEnc8;

   function KeySetupEnc(cipherKey: key; Nk: Integer) return key_schedule
   is
      rk: key_schedule;
   begin
      rk := key_schedule'(others => word'(0, 0, 0, 0));

      if (Nk = 4) then
         rk := KeySetupEnc4(cipherKey);
      end if;

      if (Nk = 6) then
         rk := KeySetupEnc6(cipherKey);
      end if;

      if (Nk = 8) then
         rk := KeySetupEnc8(cipherKey);
      end if;

      return rk;
   end KeySetupEnc;

   function KeyScheduleMod1(W: key_schedule; Nr: Integer) return key_schedule
   is
      rk: key_schedule;
   begin
      rk := W;
      for i in Integer range 0 .. Nr loop
         rk(4*i  ) := W(4*(Nr-i)  );
         rk(4*i+1) := W(4*(Nr-i)+1);
         rk(4*i+2) := W(4*(Nr-i)+2);
         rk(4*i+3) := W(4*(Nr-i)+3);
         --# assert 0 <= i and i <= Nr and Nr = Nr% and
         --#        (for all p in Integer range 0 .. i =>
         --#          (rk(4*p  ) = W(4*(Nr-p)  ) and
         --#           rk(4*p+1) = W(4*(Nr-p)+1) and
         --#           rk(4*p+2) = W(4*(Nr-p)+2) and
         --#           rk(4*p+3) = W(4*(Nr-p)+3))) and
         --#        (for all j in Integer range 4*(Nr+1) .. 4*(MAXNR+1)-1 =>
         --#          (rk(j) = W(j)));
      end loop;
      return rk;
   end KeyScheduleMod1;

   function KeyScheduleMod2(W: key_schedule; Nr: Integer) return key_schedule
   is
      st: state;
      rk: key_schedule;
   begin
      rk := W;
      for i in Integer range 1 .. Nr-1 loop
         st := state'(W(4*i), W(4*i+1), W(4*i+2), W(4*i+3));
         st := InvMixColumns(st);
         rk(4*i) := st(0); rk(4*i+1) := st(1); rk(4*i+2) := st(2); rk(4*i+3) := st(3);
         --# assert 1 <= i and i <= Nr-1 and Nr = Nr% and
         --#        (for all p in Integer range 1 .. i =>
         --#          (rk(4*p  ) = InvMixColumns(state'(W(4*p), W(4*p+1), W(4*p+2), W(4*p+3)))(0) and
         --#           rk(4*p+1) = InvMixColumns(state'(W(4*p), W(4*p+1), W(4*p+2), W(4*p+3)))(1) and
         --#           rk(4*p+2) = InvMixColumns(state'(W(4*p), W(4*p+1), W(4*p+2), W(4*p+3)))(2) and
         --#           rk(4*p+3) = InvMixColumns(state'(W(4*p), W(4*p+1), W(4*p+2), W(4*p+3)))(3))) and
         --#        (for all j in Integer range 0 .. 3 =>
         --#          (rk(j) = W(j))) and
         --#        (for all j in Integer range 4*Nr .. 4*(MAXNR+1)-1 =>
         --#          (rk(j) = W(j)));

      end loop;
      return rk;
   end KeyScheduleMod2;

   function KeySetupDec(cipherKey: key; Nk: in Integer) return key_schedule
   is
   begin
      return KeyScheduleMod2(KeyScheduleMod1(KeySetupEnc(cipherKey, Nk), Nk + 6), Nk + 6);
   end KeySetupDec;

   procedure aesKeySetupEnc(rk: out key_schedule; cipherKey: in key; keyBits: in Integer; Nr: out Integer)
   is
   begin
      Nr := keyBits/32 + 6;
      rk := KeySetupEnc(cipherKey, keyBits/32);
   end aesKeySetupEnc;

   procedure aesKeySetupDec(rk: out key_schedule; cipherKey: in key; keyBits: in Integer; Nr: out Integer)
   is
   begin
      Nr := keyBits/32 + 6;
      rk := KeySetupDec(cipherKey, keyBits/32);
   end aesKeySetupDec;

   procedure aesEncrypt(rk: in key_schedule; Nr: in Integer; pt: in block; ct: out block)
   is
      st: state;
   begin
      st := AddRoundKey(Block2State(pt), rk(0), rk(1), rk(2), rk(3));
      --# assert st = encrypt_round(Block2State(pt), rk, Nr, 0);

      for r in Natural range 1 .. Nr-1 loop
         st := AddRoundKey(MixColumns(ShiftRows(SubBytes(st))), rk(4*r), rk(4*r+1), rk(4*r+2), rk(4*r+3));
         --# assert 1 <= r and r <= Nr-1 and Nr = Nr% and
         --#        st = encrypt_round(Block2State(pt), rk, Nr, r);
      end loop;

      st := AddRoundKey(ShiftRows(SubBytes(st)), rk(4*Nr), rk(4*Nr+1), rk(4*Nr+2), rk(4*Nr+3));
      --# assert st = encrypt_round(Block2State(pt), rk, Nr, Nr);

      ct := State2Block(st);
   end aesEncrypt;

   procedure aesDecrypt(rk: in key_schedule; Nr: in Integer; ct: in block; pt: out block)
   is
      st: state;
   begin
      st := AddRoundKey(Block2State(ct), rk(0), rk(1), rk(2), rk(3));
      --# assert st = decrypt_round(Block2State(ct), rk, Nr, 0);

      for r in Natural range 1 .. Nr-1 loop
         st := AddRoundKey(InvMixColumns(InvShiftRows(InvSubBytes(st))), rk(4*r), rk(4*r+1), rk(4*r+2), rk(4*r+3));
         --# assert 1 <= r and r <= Nr-1 and Nr = Nr% and
         --#        st = decrypt_round(Block2State(ct), rk, Nr, r);
      end loop;

      st := AddRoundKey(InvShiftRows(InvSubBytes(st)), rk(4*Nr), rk(4*Nr+1), rk(4*Nr+2), rk(4*Nr+3));
      --# assert st = decrypt_round(Block2State(ct), rk, Nr, Nr);

      pt := State2Block(st);
   end aesDecrypt;
end AES;


